"""
Author: Xinyuan Ye
Description: A* Search Algorithm Visualization with automatic final path backtracking.
             Includes step-by-step exploration, and heuristic/cost display in each cell.
             The path backtracks automatically from the goal to the start after search is finished.
"""

import pygame
import heapq
import sys

# Define colors
WALKABLE = (255, 182, 193)  # Pink for walkable paths
WALL = (0, 0, 255)  # Blue for walls
START = (0, 124, 255)  # Blue for start
EXIT = (255, 0, 0)  # Red for exit
CLOSED_SET = (128, 0, 128)  # Purple for closed set (visited)
OPEN_SET = (144, 238, 144)  # Light green for open set (frontier)
PATH = (0, 255, 0)  # Green for final path
HEURISTIC_TEXT_COLOR = (0, 0, 0)  # Black for heuristic text

# Initialize pygame
pygame.init()

# Set screen size
WIDTH, HEIGHT = 600, 400
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("A* Search Visualization")

# Set font for displaying heuristic and cost values
font = pygame.font.SysFont(None, 24)

# Helper to calculate cell size
def calculate_cell_size(maze):
    cell_width = WIDTH // len(maze[0])
    cell_height = HEIGHT // len(maze)
    return cell_width, cell_height

# Load maze from a file
def load_maze_from_file(filepath):
    with open(filepath, 'r') as f:
        maze = [list(line.strip()) for line in f.readlines()]
    return maze

# Draw the current state of the maze, with the option to retain info during backtracking
def draw_maze(maze, cell_width, cell_height, f_score=None, g_score=None, path_reveal=False):
    screen.fill((0, 0, 0))  # Clear screen
    for i in range(len(maze)):
        for j in range(len(maze[0])):
            color = WALKABLE if maze[i][j] == '.' else WALL
            pygame.draw.rect(screen, color, (j * cell_width, i * cell_height, cell_width, cell_height))
            if maze[i][j] == 'S':
                pygame.draw.rect(screen, START, (j * cell_width, i * cell_height, cell_width, cell_height))
            elif maze[i][j] == 'E':
                pygame.draw.rect(screen, EXIT, (j * cell_width, i * cell_height, cell_width, cell_height))
            elif maze[i][j] == '#':
                pygame.draw.rect(screen, CLOSED_SET, (j * cell_width, i * cell_height, cell_width, cell_height))
            elif maze[i][j] == 'P':
                pygame.draw.rect(screen, PATH, (j * cell_width, i * cell_height, cell_width, cell_height))
            elif maze[i][j] == 'X':
                pygame.draw.rect(screen, OPEN_SET, (j * cell_width, i * cell_height, cell_width, cell_height))
            
            # Display heuristic and cost values (during both search and backtracking)
            if f_score and (i, j) in f_score:
                text = font.render(f"{f_score[(i, j)]}", True, HEURISTIC_TEXT_COLOR)
                screen.blit(text, (j * cell_width + cell_width // 4, i * cell_height + cell_height // 4))
            if g_score and (i, j) in g_score:
                g_text = font.render(f"G:{g_score[(i, j)]}", True, HEURISTIC_TEXT_COLOR)
                screen.blit(g_text, (j * cell_width + 5, i * cell_height + cell_height // 2))

    pygame.display.flip()

# Manhattan distance heuristic for A*
def heuristic(a, b):
    return abs(a[0] - b[0]) + abs(a[1] - b[1])

# A* Search Algorithm
def a_star(maze, start, goal):
    open_set = []
    heapq.heappush(open_set, (0, start))
    came_from = {}
    g_score = {start: 0}
    f_score = {start: heuristic(start, goal)}
    steps = []  # Store each state for replay (next/previous step)
    
    while open_set:
        pygame.event.pump()  # Handle events to avoid freezing

        current = heapq.heappop(open_set)[1]

        if current == goal:
            steps.append((maze.copy(), g_score.copy(), f_score.copy()))
            return reconstruct_path(came_from, current), steps

        x, y = current
        maze[x][y] = '#'

        # Explore neighbors (up, down, left, right)
        for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
            neighbor = (x + dx, y + dy)
            nx, ny = neighbor

            if 0 <= nx < len(maze) and 0 <= ny < len(maze[0]) and maze[nx][ny] != '*' and maze[nx][ny] != '#':
                tentative_g_score = g_score[current] + 1
                if neighbor not in g_score or tentative_g_score < g_score[neighbor]:
                    came_from[neighbor] = current
                    g_score[neighbor] = tentative_g_score
                    f_score[neighbor] = g_score[neighbor] + heuristic(neighbor, goal)
                    heapq.heappush(open_set, (f_score[neighbor], neighbor))
                    maze[nx][ny] = 'X'

        steps.append((maze.copy(), g_score.copy(), f_score.copy()))
        draw_maze(maze, cell_width, cell_height, f_score, g_score)
        pygame.time.delay(50)

    return None, steps  # No path found

# Reconstruct path after A* search completes
def reconstruct_path(came_from, current):
    path = []
    while current in came_from:
        path.append(current)
        current = came_from[current]
    path.reverse()  # From goal to start
    return path

# Automatically backtrack and visualize the final path from goal to start
def visualize_final_path(maze, path, g_score, f_score):
    for (x, y) in reversed(path):  # Backtrack from goal to start
        maze[x][y] = 'P'
        draw_maze(maze, cell_width, cell_height, f_score=f_score, g_score=g_score, path_reveal=True)
        pygame.time.delay(100)

# Main function
def main():
    # Load maze from file
    if len(sys.argv) < 2:
        print("Usage: python a_star_visualization.py <maze_file.txt>")
        sys.exit(1)
    
    maze_file = sys.argv[1]
    maze = load_maze_from_file(maze_file)

    global cell_width, cell_height
    cell_width, cell_height = calculate_cell_size(maze)

    # Find start (S) and goal (E)
    start = None
    goal = None
    for i in range(len(maze)):
        for j in range(len(maze[0])):
            if maze[i][j] == 'S':
                start = (i, j)
            elif maze[i][j] == 'E':
                goal = (i, j)

    if not start or not goal:
        print("Start or goal not defined in the maze!")
        sys.exit(1)

    # Run A* search and get the steps
    path, steps = a_star(maze, start, goal)

    running = True
    current_step = 0
    path_found = False  # Tracks whether path has been found and backtracked
    final_g_score = {}
    final_f_score = {}

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        # Automatically handle path exploration animation
        if current_step < len(steps):
            maze_state, g_score, f_score = steps[current_step]
            draw_maze(maze_state, cell_width, cell_height, f_score, g_score)
            final_g_score = g_score  # Store for backtracking phase
            final_f_score = f_score  # Store for backtracking phase
            current_step += 1
            pygame.time.delay(50)

        # Once path exploration is complete, automatically visualize final path
        if current_step == len(steps) and not path_found:
            path_found = True
            visualize_final_path(maze, path, final_g_score, final_f_score)

        pygame.time.Clock().tick(60)

    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()